import torch
from typing import Callable, TypeVar
from learners.baseline import Learner as BaselineLearner
from models.multi_net import BaseMulti

num_workers = 8

T = TypeVar("T", bound=BaseMulti)


class Learner(BaselineLearner[T]):
    """Multi-strategy learner with configurable training strategies"""

    _network: T

    def __init__(self, args, data_manager, model_func: Callable = BaseMulti):
        super().__init__(args, data_manager, model_func)

        assert args["use_proto"] in [True, False]
        self.use_proto = args["use_proto"]
        self.forward_train = lambda model, inputs: model(inputs, mode="cur")
        self.forward_test = lambda model, inputs: model(inputs, mode="all")

    def extract_token(self, inputs):
        """Only work with mode CUR, 兼容DataParallel"""
        return self._call_network_method("extract_token", inputs, mode="cur")[:, -1]

    def after_train(self):
        super().after_train()
        if self.use_proto:
            print("\n" + "=" * 60)
            print("Update fc weights with Prototypes")
            print("=" * 60 + "\n")
            self.replace_fc()

    def replace_fc(self):
        model = self._network
        model = model.eval()

        index_range = list(range(self._known_classes, self._total_classes))
        cls_protos = self._class_means[index_range]
        model.fc.heads[-1].weight.data = cls_protos.to(
            dtype=torch.float32, device=self._device
        )
